import os
import re
from string import *

import paddle
import paddlehub as hub
from PIL import Image, ImageDraw, ImageFont

'''
图片重命名、去除背景、添加图片名称进图片
'''


def rename(file_path):
    for dirpath, dirnames, filenames in os.walk(file_path):
        for filename in filenames:
            try:
                inputfilename = dirpath.replace("\\", "/") + '/' + filename
                tables = str.maketrans('', '', "()【】-  */" + r"""!"#$%&'()*+，,-/:;<=>?@[\]^_`{|}~""" + whitespace)
                # 去除中文
                pattern = re.compile(u"[\u4e00-\u9fa5]+")
                filename = re.sub(pattern, "", filename)
                outfilename = dirpath.replace("\\", "/") + '/' + str.strip(filename.translate(tables))
                print(inputfilename)
                print(outfilename)
                os.rename(inputfilename, outfilename)
            except FileExistsError as f:
                print(f)
                continue


def deleteBg(path, path_output, flag):
    paddle.enable_static()
    # 加载模型
    humanseg = hub.Module(name='deeplabv3p_xception65_humanseg')
    # 获取文件列表
    for dirpath, dirnames, filenames in os.walk(path):
        if dirpath.__contains__("changed_img"):
            continue
        for filename in filenames:
            output_dir = os.path.join(path_output, dirpath.split('\\')[-1])
            if '1'.__eq__(flag):
                output_dir = output_dir + "-h"
            if '2'.__eq__(flag):
                output_dir = output_dir + "-w"
            if '3'.__eq__(flag):
                output_dir = output_dir + "-all"
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
                print('make output_dir:' + output_dir)
            try:
                if filename.__contains__(".jpg") or filename.__contains__(".png"):
                    imgpath = os.path.join(dirpath, filename)
                    img = Image.open(imgpath)
                    # 复制高比宽分辨率大的图片
                    if '1'.__eq__(flag) and img.size[0] >= img.size[1]:
                        continue
                    if '2'.__eq__(flag) and img.size[0] < img.size[1]:
                        continue
                        # 抠图
                    results = humanseg.segmentation(data={'image': [imgpath]},
                                                    output_dir=output_dir)
                    for result in results:
                        print(result)
            except Exception as e:
                print(e)
                continue


# 将文件名添加到图片中
def add_filename_to_image(path):
    if not os.path.exists(os.path.join(path, "changed_img")):
        os.makedirs(os.path.join(path, "changed_img"))
    # 获取文件列表
    for dirpath, dirnames, filenames in os.walk(path):
        if dirpath.__contains__("changed_img"):
            continue
        for filename in filenames:
            try:
                if filename.__contains__(".jpg") or filename.__contains__(".png"):
                    imgpath = os.path.join(dirpath, filename)
                    new_imgpath = os.path.join(dirpath, "changed_img", filename)
                    print(imgpath)
                    img = Image.open(imgpath)
                    # 复制高比宽分辨率大的图片
                    # 新建绘图对象
                    draw = ImageDraw.Draw(img)
                    color = (52, 51, 85)
                    font = ImageFont.truetype("C:\Windows\Fonts\msyh.ttc", int(0.01 * img.size[0]))
                    text = '\n'.join(list(filename))
                    # print(text)
                    # 显示图片
                    draw.text((30, img.size[0] / 10), text, color=color, font=font, direction=None)
                    # 保存
                    img.save(new_imgpath)
            except Exception as e:
                print(e)
                continue


if __name__ == "__main__":
    # path = r'D:\HumanBehaviorDetectionVideo\handwashingimgused'
    flag = "2"  # 1 竖屏图片 2 横屏图片 3 全部图片
    path = r'D:\like\img-weight'
    path_output = r'D:\like\img-weight-bgdel'
    rename(path)
    deleteBg(path, path_output, flag)
    # add_filename_to_image(path)
